from selenium import webdriver
from selenium.webdriver.firefox.service import Service
import json
import os
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import pickle
import time

def save_cookies(cookies):
    """保存cookies到文件"""
    try:
        with open('xiaohongshu_cookies.pkl', 'wb') as f:
            pickle.dump(cookies, f)
        print("Cookies保存成功！")
    except Exception as e:
        print(f"保存Cookies时出错: {e}")

def load_cookies():
    """从文件加载cookies"""
    try:
        if os.path.exists('xiaohongshu_cookies.pkl'):
            with open('xiaohongshu_cookies.pkl', 'rb') as f:
                return pickle.load(f)
    except Exception as e:
        print(f"加载Cookies时出错: {e}")
    return None

def is_logged_in():
    """检查是否已登录"""
    try:
        # 尝试查找登录后才有的元素
        WebDriverWait(driver, 5).until(
            EC.presence_of_element_located((By.XPATH, '//div[contains(@class, "user-avatar")]'))
        )
        return True
    except:
        return False

def login_with_cookies():
    """使用cookies登录"""
    cookies = load_cookies()
    if cookies:
        try:
            for cookie in cookies:
                driver.add_cookie(cookie)
            driver.refresh()
            driver.implicitly_wait(10)
            
            if is_logged_in():
                print("使用cookies登录成功！")
                return True
            else:
                print("Cookies已过期，需要重新登录")
                return False
        except Exception as e:
            print(f"使用cookies登录时出错: {e}")
            return False
    return False

def phone_login():
    # 等待页面加载完成
    driver.implicitly_wait(10)
    
    # 等待手机号输入框出现并可交互
    phone_input = WebDriverWait(driver, 10).until(
        EC.element_to_be_clickable((By.XPATH, '/html/body/div[2]/div[1]/div/div[1]/div[3]/div[2]/form/label[1]/input'))
    )
    
    # 滚动到元素位置
    driver.execute_script("arguments[0].scrollIntoView(true);", phone_input)
    time.sleep(1)
    
    # 输入手机号
    phone_input.send_keys('13145037933')
    
    # 等待获取验证码按钮可点击
    verify_button = WebDriverWait(driver, 10).until(
        EC.element_to_be_clickable((By.XPATH, '/html/body/div[2]/div[1]/div/div[1]/div[3]/div[2]/form/label[2]/span'))
    )
    verify_button.click()
    
    # 等待用户手动输入验证码
    verification_code = input("请输入收到的验证码: ")
    
    # 等待验证码输入框出现并可交互
    code_input = WebDriverWait(driver, 10).until(
        EC.element_to_be_clickable((By.XPATH, '/html/body/div[2]/div[1]/div/div[1]/div[3]/div[2]/form/label[2]/input'))
    )
    code_input.send_keys(verification_code)
    
    # 等待同意政策协议按钮出现并可点击
    agree_button = WebDriverWait(driver, 10).until(
        EC.element_to_be_clickable((By.XPATH, '/html/body/div[2]/div[1]/div/div[1]/div[3]/div[3]/span/div'))
    )
    agree_button.click()
    
    # 等待登录按钮可点击
    login_button = WebDriverWait(driver, 10).until(
        EC.element_to_be_clickable((By.XPATH, '/html/body/div[2]/div[1]/div/div[1]/div[3]/div[2]/form/button'))
    )
    login_button.click()
    
    # 等待登录完成
    driver.implicitly_wait(10)
    
    # 保存cookies
    if is_logged_in():
        save_cookies(driver.get_cookies())
    
    # 等待确定按钮可点击
    confirm_button = WebDriverWait(driver, 10).until(
        EC.element_to_be_clickable((By.XPATH, '/html/body/div[2]/div[1]/div/div[1]/div[3]/div[2]/form/button'))
    )
    confirm_button.click()

def search_content():
    # 等待搜索框出现
    search_input = WebDriverWait(driver, 10).until(
        EC.presence_of_element_located((By.XPATH, '//*[@id="search-input"]'))
    )
    
    # 弹出输入框让用户输入搜索内容
    search_text = input("请输入要搜索的内容: ")
    
    # 在搜索框中输入内容
    search_input.send_keys(search_text)
    
    # 点击搜索按钮
    search_button = driver.find_element('xpath', '/html/body/div[2]/div[1]/div[1]/header/div[1]/div/div[2]')
    search_button.click()
    
    # 等待搜索结果加载
    driver.implicitly_wait(10)

def click_company():
    # 等待大厂按钮出现并可点击
    company_button = WebDriverWait(driver, 10).until(
        EC.element_to_be_clickable((By.XPATH, '/html/body/div[2]/div[1]/div[2]/div[2]/div/div/div[2]/div[2]/div/div/div[2]/button[3]'))
    )
    company_button.click()
    
    # 等待页面加载
    driver.implicitly_wait(10)
    
def get_content():
    # 等待内容加载
    driver.implicitly_wait(10)
    
    # 获取所有标题
    titles = driver.find_elements('xpath', '/html/body/div[2]/div[1]/div[2]/div[2]/div/div/div[3]/div[1]/section[1]/div/div/a/span')
    # 获取所有标题和作者
    titles = []
    authors = []
    
    # 遍历所有section获取标题和作者
    sections = driver.find_elements('xpath', '/html/body/div[2]/div[1]/div[2]/div[2]/div/div/div[3]/div[1]/section')
    
    for section in sections:
        try:
            title = section.find_element('xpath', './/div/div/a/span').text
            author = section.find_element('xpath', './/div/div/div/a/span').text
            titles.append(title)
            authors.append(author)
        except:
            continue
    
    # 打印所有标题和作者
    print("\n=== 搜索结果 ===")
    for title, author in zip(titles, authors):
        print(f"标题: {title}")
        print(f"作者: {author}")
        print("-" * 50)
    # 获取所有作者名称
    authors = driver.find_elements('xpath', '/html/body/div[2]/div[1]/div[2]/div[2]/div/div/div[3]/div[1]/section[1]/div/div/div/a/span')
    
    # 打印标题和作者
    print("\n=== 搜索结果 ===")
    for title, author in zip(titles, authors):
        print(f"标题: {title.text}")
        print(f"作者: {author.text}")
        print("-" * 50)

# 请替换为你实际的 GeckoDriver 路径
geckodriver_path = '/Users/lidongxiao/Downloads/geckodriver'
service = Service(geckodriver_path)

driver = webdriver.Firefox(service=service)
try:
    driver.get('https://www.xiaohongshu.com/explore')
    # 等待页面加载
    driver.implicitly_wait(10)
    
    # 尝试使用cookies登录
    if not login_with_cookies():
        print("使用cookies登录失败，尝试使用手机号登录...")
        phone_login()
    
    # 执行搜索功能
    search_content()
    
    # 点击大厂按钮
    click_company()
    
    # 获取内容
    get_content()
    
    input("按回车键继续...")
finally:
    driver.quit()